{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TypeApplications #-}
{-# LANGUAGE TypeInType #-}
{-# LANGUAGE TypeOperators #-}

-- | Module: Data.Vector.Mutable.Sized.PartialSort
-- Copyright: (c) 2018 Koz Ross
-- Maintainer: Koz Ross <koz.ross@retro-freedom.nz>
-- Stability: Experimental
-- Portability: Non-portable
--
-- Partial sorts for mutable vectors, using heapsort with a quarternary heap. These are performed in-place.
module Data.Vector.Mutable.Sized.PartialSort
  ( partialSort
  , partialSortBy
  , partialSortWithin
  , partialSortByWithin
  ) where

import Control.Monad.Primitive
import Data.Coerce
import Data.Finite
import Data.Functor.Contravariant
import GHC.TypeNats

import qualified Data.Vector.Algorithms.Heap as H
import qualified Data.Vector.Generic.Mutable as VGM
import qualified Data.Vector.Generic.Mutable.Sized as VGMS

-- | Partial sorting. After @partialSort v k@, the first @k@ elements of @v@ will be in order, with
-- the rest in an unspecified order.
--
-- __Complexity:__ \(O(n \log (n))\)
{-# INLINE partialSort #-}
partialSort ::
     (Ord a, PrimMonad m, VGM.MVector v a)
  => VGMS.MVector v n (PrimState m) a
  -> Finite (n + 1)
  -> m ()
partialSort v = H.partialSort (VGMS.fromSized v) . unpack

-- | As 'partialSort', but with a custom 'Comparison'.
{-# INLINE partialSortBy #-}
partialSortBy ::
     (PrimMonad m, VGM.MVector v a)
  => Comparison a
  -> VGMS.MVector v n (PrimState m) a
  -> Finite (n + 1)
  -> m ()
partialSortBy comp v = H.partialSortBy (coerce comp) (VGMS.fromSized v) . unpack

-- | Partially sorts within the specified indices.
-- More precisely, @partialSortWithin v i j@ performs a partial sort of all
-- elements of @v@ between @min i j@ (inclusive) and @max i j@ (exclusive). In
-- particular, when @i == j@, no sorting will happen. Elements outside of this
-- index range will be in no particular order. 
{-# INLINE partialSortWithin #-}
partialSortWithin ::
     forall m v n a. (PrimMonad m, VGM.MVector v a, Ord a)
  => VGMS.MVector v n (PrimState m) a
  -> Finite n
  -> Finite n
  -> m ()
partialSortWithin v i j =
  H.partialSortByBounds
    (coerce (defaultComparison @a))
    (VGMS.fromSized v)
    (abs (i' - j'))
    (min i' j')
    (max i' j')
  where
    i' = unpack i
    j' = unpack j

-- | As 'partialSortWithin', but with a custom 'Comparison'.
{-# INLINE partialSortByWithin #-}
partialSortByWithin ::
     (PrimMonad m, VGM.MVector v a)
  => Comparison a
  -> VGMS.MVector v n (PrimState m) a
  -> Finite n
  -> Finite n
  -> m ()
partialSortByWithin comp v i j =
  H.partialSortByBounds
    (coerce comp)
    (VGMS.fromSized v)
    (abs (i' - j'))
    (min i' j')
    (max i' j')
  where
    i' = unpack i
    j' = unpack j

-- helper for Finite unpacking
{-# INLINE unpack #-}
unpack :: Finite n -> Int
unpack = fromIntegral . getFinite
